package sdp.server.vision;

import java.awt.image.BufferedImage;

import sdp.server.BallState;
import sdp.server.RobotState;
import sdp.server.ScreenProjection;
import sdp.server.math.Vector2;

import static com.googlecode.javacv.cpp.opencv_core.*;
import static com.googlecode.javacv.cpp.opencv_imgproc.*;


public class Vision {
	public static final int BALL_CHANNEL         = 0;
	public static final int YELLOW_CHANNEL       = 1;
	public static final int BLUE_CHANNEL         = 2;
	public static final int YELLOW_BLACK_CHANNEL = 3;
	public static final int BLUE_BLACK_CHANNEL   = 4;
	public static final int N_ENTITY_CHANNELS    = 5;
	public static final int N_ENTITY_MODELS      = 3;
		
	private static String dirRoot    = ".";
	private static String dataFolder = dirRoot + "/Data";
	
	private CvMemStorage storage = CvMemStorage.create();
	
	private IplImage         raw, undistorted, cropped;
	private IplImage[]       entityChannels;
	private IplImage         temp1c, temp3c, temp3u;
	private IplImage[]       hsv;
	private IplImage         mapx, mapy;
	private IplImage         debugImage;
	private ScreenProjection screen; 
	
	private int frameId     = 0;
	private int showChannel = -1;
	
	private CvScalar[] entDbgCols = { cvScalar(0,0,255,0), cvScalar(0,255,255,0), cvScalar(255,0,0,0) };
	
	private BallModel     ballModel;
	private RobotModel[]  robotModels;
	private EntityModel[] entityModels;
	
	
	// PITCH 1
	private boolean performUndistortion = false;
	private CvRect  cropRect            = cvRect(0, 130, 700, 395);
	private String  entityPropsFile     = "entity_props_pitch1.xml";
	
	// PITCH 2
	//private boolean performUndistortion = true;
	//private CvRect  cropRect            = cvRect(30, 105, 640, 360);
	//private String  entityPropsFile     = "entity_props_pitch2.xml";
	
	
	
	public enum Value {
		MEAN,
		RANGE
	}
	
	public enum Property {
		HUE, SATURATION, VALUE
	}
	
	public Vision( ScreenProjection screenProjection ) {
		screen         = screenProjection;
		entityModels   = new EntityModel[ N_ENTITY_MODELS ];
		entityChannels = new IplImage[ N_ENTITY_CHANNELS ];
		robotModels    = new RobotModel[ 2 ];
		hsv            = new IplImage[ 3 ];
		updateProjection();
	}
	
	public ScreenProjection getScreenProjection() {
		return screen;
	}
	
	public void initialise( BufferedImage frame ) {
		allocUncroppedImages( frame );
		allocCroppedImages();
		loadCalibrationInfo();
		createEntityModels();
		loadProperties();
	}
	
	public void processImage( BufferedImage frame, RobotState[] robotStates, BallState ball, double timeStep ) {
		//long time = System.currentTimeMillis();
		storage.release();
		storage = CvMemStorage.create();
		for( EntityModel em : entityModels ) {
			em.setStorage( storage );
		}
		
		
		raw.copyFrom( frame );
		
		undistort();
		applyCrop();
		
		debugImage = detectEntities( robotStates, ball, timeStep );
		frameId ++;
		
		//System.out.println( "Vision: processImage(): Everything - " + (System.currentTimeMillis() - time) + " ms." );
	}
	
	public void updateProjection() {
		screen.setScreenSize( new Vector2( cropRect.width(), cropRect.height() ) );
		screen.setScale( ( 710 * 0.00344 ) / cropRect.width() );
	}
	
	public IplImage getDebugImage() {
		return debugImage;
	}
	
	public boolean isShowingChannel() {
		// TODO: Change channel showing into a more robust debug level system.
		return showChannel != -1;
	}
	
	public void showEntityChannel( int newChannel ) {
		showChannel = newChannel; 
		if( isShowingChannel() && showChannel < 3 ) {
			for( EntityModel em : entityModels )
				em.setDebugImage( debugImage );
		} else {
			for( EntityModel em : entityModels )
				em.setDebugImage( null );
		}
	}
	
	public void showNormal() {
		showChannel = -1;
		for( EntityModel em : entityModels )
			em.setDebugImage( null );
	}
	
	// TODO: These methods, while general, are pretty cumbersome to use. Would be nice to have
	// convenience methods that wrap them (e.g. setMeanHue( entity, newMeanValue ) etc. ).
	public void setProperty( int entity, Property prop, Value valueType, int value ) {
		EntityModel em;
		if( entity < 3 )
			em = entityModels[ entity ];
		else
			em = robotModels[ entity - 3 ].getBlackCircleModel();
		
		if( valueType == Value.MEAN )
			em.setHsvMean( prop.ordinal(), value );
		else
			em.setHsvRange( prop.ordinal(), value );
	}
	
	public int getProperty( int entity, Property prop, Value valueType ) {
		EntityModel em;
		if( entity < 3 )
			em = entityModels[ entity ];
		else
			em = robotModels[ entity - 3 ].getBlackCircleModel();
		
		if( valueType == Value.MEAN )
			return em.getHsvMean( prop.ordinal() );
		else
			return em.getHsvRange( prop.ordinal() );
	}
	
	public int getPropertyMean( int type, Property prop ) {
		return getProperty( type, prop, Value.MEAN );
	}
	
	public int getPropertyRange( int type, Property prop ) {
		return getProperty( type, prop, Value.RANGE );
	}
	
	public int getPropertyMin( int type, Property prop ) {
		return clamp255( getProperty( type, prop, Value.MEAN) -
		                 getProperty( type, prop, Value.RANGE)
		);
	}
	
	public int getPropertyMax( int type, Property prop ) {
		return clamp255( getProperty( type, prop, Value.MEAN) +
		                 getProperty( type, prop, Value.RANGE)
		);
	}
	
	public void setCropLeft( int newValue ) {
		cropRect.x( newValue );
	}
	
	public void setCropTop( int newValue ) {
		cropRect.y( newValue );
	}
	
	public void setCropWidth( int newValue ) {
		cropRect.width( newValue );
	}
	
	public void setCropHeight( int newValue ) {
		cropRect.height( newValue );
	}
	
	public int getCropLeft() {
		return cropRect.x();
	}
	
	public int getCropTop() {
		return cropRect.y();
	}
	
	public int getCropWidth() {
		return cropRect.width();
	}
	
	public int getCropHeight() {
		return cropRect.height();
	}
	
	public void saveProperties() {
		saveProperties( entityPropsFile );
	}
	
	public void saveProperties( String filename ) {
		CvMat data = new CvMat( cvLoad(dataFolder + "/" + filename) );
		
		if( !data.isNull() ) {
			System.out.println("Vision: Property file already exists, backing up...");
			cvSave(dataFolder + "/b4save_" + filename, data);
		}
		
		int m = Property.values().length;
		int n = N_ENTITY_CHANNELS;
		int p = Value.values().length;
		
		data = CvMat.create( m*n*p, 1 );
		
		int s = 0;
		for( int i = 0 ; i < m ; ++ i ) {
			for( int j = 0 ; j < n ; ++ j ) {
				if( j < 3 ) {
					data.put(s ++, entityModels[ j ].getHsvMean( i ) );
					data.put(s ++, entityModels[ j ].getHsvRange( i ) );
				} else {
					data.put(s ++, robotModels[ j - 3 ].getBlackCircleModel().getHsvMean( i ) );
					data.put(s ++, robotModels[ j - 3 ].getBlackCircleModel().getHsvRange( i ) );
				}
				
			}
		}
		
		cvSave(dataFolder + "/" + filename, data);
		System.out.println("Vision: Properties saved.");
	}
	
	public void loadProperties() {
		loadProperties( entityPropsFile );
	}
	
	public void loadProperties( String filename ) {
		System.out.println("Vision: Saving current properties...");
		saveProperties( "b4load_entity_props.xml" );
		System.out.println("Vision: Loading new properties...");
		CvMat data = new CvMat( cvLoad(dataFolder + "/" + filename) );
		
		int m = Property.values().length;
		int n = 5;
		int p = Value.values().length;
		
		if( data.isNull() ) {
			System.err.println("Vision: Could not open file.");
		} else if( data.rows() != m*n*p || data.cols() != 1 ) {
			System.err.println("Vision: Invalid dimensions in file.");
		} else {
			int s = 0;
			for( int i = 0 ; i < m ; ++ i ) {
				for( int j = 0 ; j < n ; ++ j ) {
					if( j < 3 ) {
						entityModels[ j ].setHsvMean( i, (int)data.get( s ++ ) );
						entityModels[ j ].setHsvRange( i, (int)data.get( s ++ ) );
					} else {
						robotModels[ j - 3 ].getBlackCircleModel().setHsvMean( i, (int)data.get( s ++ ) );
						robotModels[ j - 3 ].getBlackCircleModel().setHsvRange( i, (int)data.get( s ++ ) );
					}
				}
			}
		}
	}
	
	public BufferedImage autoCrop() {
		/*
		IplImage blur  = uncropped3[ 0 ], hsv = uncropped3[ 1 ];
		IplImage hue   = uncropped1[ 0 ], sat = uncropped1[ 1 ], val = uncropped1[ 2 ];
		IplImage tmp1  = uncropped1[ 3 ], tmp2 = uncropped1[ 4 ];
		IplImage pitch = uncropped1[ 5 ];
		
		// Add a linear blur filter to get rid of a bit of grain (we don't care about fine detail here)
		cvSmooth( raw, blur, CV_BLUR, 5 );
		
		// Convert the blurred image to HSV, save the channels in hue, sat and val
		cvCvtColor( blur, hsv, CV_BGR2HSV );
		cvSplit( hsv, hue, sat, val, null );
		
		// Select pixels with minHue[PITCH] < hue < maxHue[PITCH], save mask in pitch
		cvThreshold( hue, tmp1, getPropertyMin( Entity.PITCH, Property.HUE ), 255, CV_THRESH_BINARY );
		cvThreshold( hue, tmp2, getPropertyMax( Entity.PITCH, Property.HUE ), 255, CV_THRESH_BINARY_INV );
		cvAnd( tmp1, tmp2, pitch, null );
		
		// Further select non-dark pixels (val > minValue / 2) to get rid of black boundaries of the pitch.
		cvThreshold( val, tmp1, getPropertyMin( Entity.PITCH, Property.VALUE ) * .5, 255, CV_THRESH_BINARY );
		cvAnd( pitch, tmp1, pitch, null );
		
		// Get the contours in the pitch mask
		CvSeq seq = new CvSeq(null);
		cvFindContours(pitch, storage, seq,Loader.sizeof(CvContour.class), CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE);
		
		// If there are any contours, find the one with the largest bounding box in pitch
		if( !seq.isNull() ) {
			CvRect maxRect = null;
			for( CvSeq s = seq ; s != null && !s.isNull() ; s = s.h_next() ) {
				CvRect r = cvBoundingRect(s, 0);
				if( maxRect == null || r.width()*r.height() > maxRect.width()*maxRect.height() )
					maxRect = r;
			}
			
			// Update the crop rectangle to the max rectangle
			System.out.println("Vision: Pitch: Rect = " + maxRect.x() + ", " + maxRect.y() + " ; " + maxRect.width() + ", " + maxRect.height() );
			cropRect = maxRect;
			
			// Realloc images to the size of the crop rectangle
			allocCroppedImages();
			
			// Ensure the projection uses the new screen size
			updateProjection();
		}
		
		return raw.getBufferedImage();
		*/
		return null;
	}
	
	private IplImage detectEntities( RobotState[] robots, BallState ball, double timeStep ) {
		cvSmooth( cropped, temp3c, CV_GAUSSIAN, 3 );
		
		// Separate image into hue, sat and value
		cvCvtColor( temp3c, temp3c, CV_BGR2HSV );
		cvSplit( temp3c, hsv[0], hsv[1], hsv[2], null );
		
		// We will be drawing debug info on the image, so we want to leave
		// cropped unchanged.
		cvCopy( cropped, debugImage );
		
		// Threshold the HSV for each of the entity models.
		for( EntityModel em : entityModels )
			em.threshold( hsv );
		
		
		if( isShowingChannel() ) {
			if( showChannel < 3 ) {
				// If we're debugging one of the object channels, then we show all of them.
				cvSet( temp1c, cvScalarAll(0) );
				for( int c = 0 ; c < 3 ; ++ c )
					cvOr( temp1c, entityChannels[ c ], temp1c, null);
				 
				cvMerge( entityChannels[2], entityChannels[1], entityChannels[0], null, temp3c );
				cvAndS( debugImage, cvScalarAll(0), debugImage, temp1c );
				cvOr( debugImage, temp3c, debugImage, null );
				 
				for( int i = 0 ; i < 3 ; ++ i )
					cvRectangleR( debugImage, entityModels[ i ].getRoi(), cvScalar(255,255,255,0), 1, 8, 0);
				
				for( int i = 0 ; i < 2 ; ++ i )
					cvRectangleR( debugImage, robotModels[ i ].getBlackCircleModel().getRoi(), cvScalar(255,255,255,0), 1, 8, 0);
				
			} else {
				// Otherwise, if we're debugging the pitch or black circle channel, show only that.
				cvMerge(
				        entityChannels[ showChannel ],
				        entityChannels[ showChannel ],
				        entityChannels[ showChannel ],
				        null, debugImage
				);
			}
		}
		
		for( EntityModel em : entityModels )
			em.detect();
		
		for( int i = 0 ; i < 2 ; ++ i )
			robotModels[ i ].update( robots[ i ], timeStep );
		
		ballModel.update( ball, timeStep );
		
		if( !isShowingChannel() ) 
			drawDetectedEntites( debugImage, robots, ball );
		
		return debugImage;
	}
	
	private void drawDetectedEntites( IplImage image, RobotState[] robots, BallState ball ) {
		// Draw robots
		for( RobotState robot : robots ) {
			if( robot.isDetected() ) {
				Vector2 p1 = screen.projectPosition( robot.getPosition() );
				Vector2 p2 = p1.plus( Vector2.fromAngle( -robot.getRotation() ).times( 25.0 ) );
				CvPoint cp1 = p1.getCvPoint(), cp2 = p2.getCvPoint();
				CvScalar robotColour = entDbgCols[ robot.getTeam().ordinal() + 1 ];
				cvLine( image, cp1, cp2, cvScalarAll(0), 3, 8, 0 );
				cvLine( image, cp1, cp2, robotColour, 2, 8, 0 );
			}
		}
		
		// Draw ball
		drawCenteredRectangle(
				image,
				screen.projectPosition( ball.getPosition() ),
				10,
				ball.isDetected() ? cvScalar(255, 255, 0, 0) : cvScalar( 0, 0, 255, 0 ),
				2
		);
			
		Vector2 p1 = ball.getPosition();
		Vector2 p2 = p1.plus( ball.getLinearVelocity().times(1.0/12.0) );
		
		p1 = screen.projectPosition( p1 ); p2 = screen.projectPosition( p2 );
		cvLine( image, p1.getCvPoint(), p2.getCvPoint(), cvScalarAll(255), 1, 8, 0 );
	}
	
	public static void drawCenteredRectangle( IplImage image, Vector2 p, int size, CvScalar colour, int thickness ) {
		Vector2 szv = new Vector2(size/2.0,size/2.0);
		CvPoint p1 = p.minus( szv ).getCvPoint();
		CvPoint p2 = p.plus( szv ).getCvPoint();
		cvRectangle( image, p1, p2, colour, thickness, 8, 0);
	}

	private void loadCalibrationInfo() {
		CvMat intrinsics = new CvMat(cvLoad(dataFolder + "/Intrinsics.yml"));
		CvMat distortion = new CvMat(cvLoad(dataFolder + "/Distortion.yml"));
		
		if( intrinsics == null || distortion == null ) {
			System.err.println("Can't open distortion info.");
		} else {
			mapx = IplImage.create(cvGetSize(raw), IPL_DEPTH_32F, 1);
			mapy = IplImage.create(cvGetSize(raw), IPL_DEPTH_32F, 1);
			cvInitUndistortMap(intrinsics, distortion, mapx, mapy);
		}
	}
	
	private void undistort() {
		if( mapx == null || mapy == null || !performUndistortion ) {
			undistorted = raw;
		} else {
			cvCopy( raw, temp3u );
			cvRemap( temp3u, undistorted, mapx, mapy, CV_INTER_LINEAR | CV_WARP_FILL_OUTLIERS, cvScalarAll(0) );
		}
	}
	
	private void applyCrop() {
		clampCropRect( cropRect, undistorted );
		cvSetImageROI( undistorted, cropRect );
		cvCopy( undistorted, cropped );
		cvResetImageROI( undistorted );
	}
	
	private void createEntityModels() {
		entityModels[ 0 ] = ballModel = new BallModel(
			storage,             // memory storage
			entityChannels[ 0 ], // ball channel
			temp1c,              // temporary image, 1 channel, cropped
			screen               // screen projection
		);
		
		for( int i = 0 ; i < 2 ; ++ i ) {
			entityModels[ i + 1 ] = robotModels[ i ] = new RobotModel(
				storage,                 // memory storage
				entityChannels[ i + 1 ], // robot channel
				entityChannels[ i + 3 ], // black circle channel
				temp1c,                  // temporary image, 1 channel, cropped
				screen                   // screen projection
			);
		}
	}
	
	private void allocUncroppedImages( BufferedImage frame ) {
		raw         = IplImage.createFrom( frame );
		undistorted = newImage( raw, 3 );
		temp3u      = newImage( raw, 3 );
	}
	
	private void allocCroppedImages() {
		assert raw != null;
		
		cropped = IplImage.create( cvSize(cropRect.width(), cropRect.height()), raw.depth(), 3 );
		temp3c  = newImage( cropped, 3 );
		temp1c  = newImage( cropped, 1 );
		
		for( int i = 0 ; i < entityChannels.length ; ++ i )
			entityChannels[ i ] = newImage( cropped, 1 );
		
		for( int i = 0 ; i < 3; ++ i )
			hsv[ i ] = newImage( cropped, 1 );
		
		debugImage = newImage( cropped, 3 );
	}
	
	private static IplImage newImage( IplImage img, int channels ) {
		return IplImage.create( cvGetSize(img), img.depth(), channels );
	}
	
	public static void clampCropRect( CvRect crop, IplImage img ) {
		if( crop.x() < 0 ) crop.x( 0 );
		if( crop.y() < 0 ) crop.y( 0 );
		
		if( crop.width()  + crop.x() > img.width() )
			crop.width( img.width() - crop.x() );
		
		if( crop.height() + crop.y() > img.height() )
			crop.height( img.height() - crop.y() );
	}
	
	private static int clamp255( int x ) {
		if( x < 0 )
			return 0;
		else if( x > 255 )
			return 255;
		else
			return x;
	}
}
